import os
import glob
import tarfile

import numpy as np
import tensorflow as tf
from configuration import ExtractionConfig

os.environ['CUDA_VISIBLE_DEVICES'] = '2'
layer_to_extract = 'pool_3:0'

config = ExtractionConfig()


def main(_):
    """提取directory中所以图片的feature，dict：{path: array}"""
    if not os.path.exists(config.img_dir):
        print("image_dir does not exit!")
        return None

    # 解压预训练model
    file_name = 'inception-2015-12-05.tgz'
    file_path = os.path.join(config.model_dir, file_name)
    tarfile.open(file_path, 'r:gz').extractall(config.pre_trained_folder)

    # import预训练的graph
    pre_trained_model_name = 'classify_image_graph_def.pb'
    with tf.gfile.FastGFile(os.path.join(config.pre_trained_folder, pre_trained_model_name), 'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
        _ = tf.import_graph_def(graph_def, name='')

    with tf.Session() as sess:
        image_path_to_vector = {}
        extract_tensor = sess.graph.get_tensor_by_name(layer_to_extract)
        counter = 0
        img_paths = glob.glob(config.img_dir + '*.jpg')
        num_images = len(img_paths)

        print("There are total " + str(num_images) + " images to process.")
        for img_idx in range(num_images):
            if config.flag > 0:
                counter += 1
                if counter % config.flag == 0:
                    print("Processing {}th image".format(img_idx))

            temp_path = img_paths[img_idx]
            image_data = tf.gfile.FastGFile(temp_path, 'rb').read()

            predictions = sess.run(extract_tensor, {'DecodeJpeg/contents:0': image_data})
            predictions = np.squeeze(predictions)

            image_path_to_vector[temp_path] = predictions

        np.save(config.save_dir, image_path_to_vector)

        print("Extracted features saved in: ", config.save_dir)


if __name__ == '__main__':
    tf.app.run()
